from typing import Optional, Callable, Any
from datetime import timedelta
import uuid
from utils.logger import Logger
from utils.exceptions import LockAcquireError
from middleware.redis_adapter import get_redis_adapter

# 初始化日志
logger = Logger.get_logger(agent_name="distributed_lock_agent")


class DistributedLock:
    """基于Redis的分布式锁工具"""

    def __init__(self):
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.lock_prefix = "distributed:lock:"  # 锁键前缀
        self.default_timeout = 30  # 默认锁超时时间（30秒，防止死锁）
        logger.info("分布式锁工具初始化成功，默认超时时间: 30秒")

    def _get_lock_key(self, lock_name: str) -> str:
        """获取锁的Redis键"""
        return f"{self.lock_prefix}{lock_name}"

    def acquire(
            self,
            lock_name: str,
            timeout: int = None,
            retry_count: int = 3,
            retry_interval: float = 0.5
    ) -> Optional[str]:
        """
        获取分布式锁
        Args:
            lock_name: 锁名称（唯一标识锁资源）
            timeout: 锁超时时间（秒），默认30秒
            retry_count: 获取失败后的重试次数
            retry_interval: 重试间隔（秒）
        Returns:
            锁标识（解锁时需要），获取失败返回None
        """
        timeout = timeout or self.default_timeout
        lock_key = self._get_lock_key(lock_name)
        lock_id = str(uuid.uuid4())  # 唯一锁标识（防止误解锁）
        retry = 0

        while retry <= retry_count:
            try:
                # Redis SET NX + EX 原子操作（不存在则设置，同时设置过期时间）
                acquired = self.redis_client.set(
                    key=lock_key,
                    value=lock_id,
                    nx=True,  # 仅当键不存在时设置
                    ex=timeout  # 过期时间（秒）
                )

                if acquired:
                    logger.debug(f"获取分布式锁成功: lock_name={lock_name}, lock_id={lock_id}, timeout={timeout}s")
                    return lock_id
                else:
                    # 锁已被占用，重试
                    retry += 1
                    logger.debug(f"获取分布式锁失败（已被占用）: lock_name={lock_name}, 重试次数={retry}/{retry_count}")
                    if retry <= retry_count:
                        import time
                        time.sleep(retry_interval)
            except Exception as e:
                logger.error(f"获取分布式锁异常: lock_name={lock_name}, error={str(e)}", exc_info=True)
                retry += 1
                if retry <= retry_count:
                    import time
                    time.sleep(retry_interval)

        logger.error(f"获取分布式锁失败（重试次数耗尽）: lock_name={lock_name}")
        return None

    def acquire_with_exception(
            self,
            lock_name: str,
            timeout: int = None,
            retry_count: int = 3,
            retry_interval: float = 0.5
    ) -> str:
        """
        获取分布式锁（失败抛出异常）
        Returns:
            锁标识（解锁时需要）
        Raises:
            LockAcquireError: 获取锁失败
        """
        lock_id = self.acquire(lock_name, timeout, retry_count, retry_interval)
        if not lock_id:
            raise LockAcquireError(
                message=f"获取分布式锁失败（重试{retry_count}次耗尽）",
                context={"lock_name": lock_name}
            )
        return lock_id

    def release(self, lock_name: str, lock_id: str) -> bool:
        """
        释放分布式锁（仅释放自己获取的锁）
        Args:
            lock_name: 锁名称
            lock_id: 锁标识（acquire返回的ID）
        Returns:
            释放成功返回True，失败返回False
        """
        lock_key = self._get_lock_key(lock_name)
        try:
            # 用Lua脚本保证原子性（判断锁标识是否匹配，匹配则删除）
            lua_script = """
            if redis.call('get', KEYS[1]) == ARGV[1] then
                return redis.call('del', KEYS[1])
            else
                return 0
            end
            """
            # 替换ARGV[1]为实际lock_id
            lua_script = lua_script.replace("ARGV[1]", f"'{lock_id}'")
            result = self.redis_client.eval(lua_script, keys=[lock_key], args=[])

            if result == 1:
                logger.debug(f"释放分布式锁成功: lock_name={lock_name}, lock_id={lock_id}")
                return True
            else:
                # 锁不存在或标识不匹配（可能已过期或被其他进程释放）
                logger.warning(f"释放分布式锁失败: lock_name={lock_name}, lock_id={lock_id}（锁不存在或标识不匹配）")
                return False
        except Exception as e:
            logger.error(f"释放分布式锁异常: lock_name={lock_name}, lock_id={lock_id}, error={str(e)}", exc_info=True)
            return False

    def with_lock(
            self,
            lock_name: str,
            timeout: int = None,
            retry_count: int = 3,
            retry_interval: float = 0.5
    ) -> Callable:
        """
        装饰器：自动获取锁和释放锁
        Usage:
            @distributed_lock.with_lock("test_lock", timeout=10)
            def func():
                pass
        """

        def decorator(func: Callable) -> Callable:
            def wrapper(*args, **kwargs) -> Any:
                lock_id = None
                try:
                    # 获取锁
                    lock_id = self.acquire_with_exception(lock_name, timeout, retry_count, retry_interval)
                    # 执行函数
                    return func(*args, **kwargs)
                finally:
                    # 释放锁（无论函数是否异常）
                    if lock_id:
                        self.release(lock_name, lock_id)

            return wrapper

        return decorator


# 单例实例（全局复用）
distributed_lock = DistributedLock()

# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="distributed_lock_test", agent_name="distributed_lock_agent")
    test_logger.info("开始测试分布式锁工具")

    # 测试1：基本获取/释放
    lock_name = "test_lock_001"
    lock_id = distributed_lock.acquire(lock_name, timeout=10)
    if lock_id:
        # 模拟业务操作
        import time

        time.sleep(2)
        # 释放锁
        distributed_lock.release(lock_name, lock_id)


    # 测试2：装饰器用法
    @distributed_lock.with_lock("test_lock_002", timeout=10)
    def test_func():
        test_logger.info("装饰器模式：获取锁后执行业务逻辑")
        time.sleep(3)


    test_func()

    # 测试3：失败抛出异常
    try:
        # 模拟并发获取锁（第一个线程占用，第二个线程失败）
        lock_id1 = distributed_lock.acquire_with_exception("test_lock_003", timeout=5)
        # 第二个线程尝试获取（重试1次）
        lock_id2 = distributed_lock.acquire_with_exception("test_lock_003", retry_count=1)
    except LockAcquireError as e:
        test_logger.error(f"并发获取锁失败: {e}")
        distributed_lock.release("test_lock_003", lock_id1)

    test_logger.success("分布式锁工具测试完成")